package com.beta.cat.config;

import org.apache.commons.lang3.StringUtils;
import org.apache.dubbo.config.spring.extension.SpringExtensionFactory;
import org.apache.dubbo.rpc.RpcContext;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.web.servlet.DispatcherServlet;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.EnumSet;

/**
 * 链路跟踪的实现
 */
@Configuration
@ConditionalOnProperty(name = "com.beta.tracelog.enable", havingValue = "true", matchIfMissing = true)
@ConditionalOnClass({LoggerFactory.class, Servlet.class, DispatcherServlet.class, WebMvcConfigurer.class, RpcContext.class})
public class BetaTraceLoggerAutoConfiguration {
    @Value("${spring.application.name:notfound}")
    private String applicationName;

    @Bean
    public FilterRegistrationBean createFilter(ITraceMDC mdc) {
        FilterRegistrationBean filterRegistration = new FilterRegistrationBean();

        TraceLogFilter filter = new TraceLogFilter(mdc);
        filterRegistration.setFilter(filter);
        filterRegistration.addUrlPatterns("/*");
        filterRegistration.setDispatcherTypes(EnumSet.allOf(DispatcherType.class));
        filterRegistration.setOrder(Ordered.HIGHEST_PRECEDENCE + 10);
        return filterRegistration;
    }

    /**
     * 日志的traceId记录
     */
    class TraceLogFilter implements Filter {
        private ITraceMDC mdc;

        TraceLogFilter(ITraceMDC mdc) {
            this.mdc = mdc;
        }
        public void init(FilterConfig filterConfig) {
        }
        public void destroy() {
        }

        @Override
        public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
            if (mdc == null || !(request instanceof HttpServletRequest) || !(response instanceof HttpServletResponse)) {
                chain.doFilter(request, response);
                return;//如果没有日志组件，则不处理
            }
            try {
                String traceId = ((HttpServletRequest) request).getHeader(RpcTraceFilter.MDC_TRACE_KEY);
                if (StringUtils.isBlank(traceId)) {
                    traceId = RpcTraceFilter.getUUID();
                }
                RpcContext.getContext().setAttachment(RpcTraceFilter.MDC_TRACE_KEY, traceId);
                ((HttpServletResponse) response).setHeader(RpcTraceFilter.MDC_TRACE_KEY, traceId);

                //向日志的Context中写入：应用名称、traceId（如果没有，则创建）
                mdc.put(RpcTraceFilter.MDC_TRACE_KEY, traceId);
                mdc.put(RpcTraceFilter.MDC_APPNAME_KEY, applicationName);

                chain.doFilter(request, response);
            } finally {
                //清空对应的内容，重新记录
                mdc.put(RpcTraceFilter.MDC_TRACE_KEY, null);
                mdc.put(RpcTraceFilter.MDC_APPNAME_KEY, null);
                RpcContext.getContext().setAttachment(RpcTraceFilter.MDC_TRACE_KEY, null);
            }
        }
    }

    @Bean
    @ConditionalOnClass(MDC.class)
    @ConditionalOnMissingBean(ITraceMDC.class)
    public ITraceMDC slf4jMDC(ApplicationContext context) {
        SpringExtensionFactory.addApplicationContext(context);
        return new ITraceMDC() {
            @Override
            public void put(String key, String val) {
                MDC.put(key, val);
            }

            @Override
            public String get(String key) {
                return MDC.get(key);
            }
        };
    }
}
